#include <occmesh/api/arap.h>

const float MIN_DEFAULT = FLT_MAX;
const float MAX_DEFAULT = -FLT_MAX;

ARAPBuilder::ARAPBuilder(Mesh *mesh, const std::vector<Chart> &charts, double lambda)
    : m_mesh(mesh), m_charts(charts), m_lambda(lambda)
{
    m_mesh->add_property(m_Lts);
    m_mesh->add_property(m_halfedgeData);
    m_mesh->add_property(m_ids);
    m_mesh->add_property(m_chartIds);
}

ARAPBuilder::~ARAPBuilder()
{
    m_mesh->remove_property(m_Lts);
    m_mesh->remove_property(m_halfedgeData);
    m_mesh->remove_property(m_ids);
    m_mesh->remove_property(m_chartIds);
}

void ARAPBuilder::Build(int times)
{
    // 标记图表的面
    for (unsigned int i = 0; i < m_charts.size(); i++)
    {
        auto &chart = m_charts[i];
        for (auto &face : chart.faces())
        {
            // 面标记为所属图表的 Id
            m_mesh->property(m_chartIds, face) = chart.id();
        }
    }

    // 开始计算
    for (unsigned int i = 0; i < m_charts.size(); i++)
    {
        // 初始化数据
        auto &chart = m_charts[i];
        SparseMatrix MT = initializeData(chart);

        int t = 0;
        while (t < times)
        {
            localPhase(chart);
            globalPhase(chart, MT);
            t++;
        }

        // 保存并规范化
        storeUVs(chart);
    }
}

SparseMatrix ARAPBuilder::initializeData(const Chart &chart)
{
    // 计算每个面的信息并保存
    auto &faces = chart.faces();
    for (auto &face : faces)
    {
        int i = 0;
        auto h_it = m_mesh->fh_ccwbegin(face), h_end = m_mesh->fh_ccwend(face);

        // 记录句柄和方向
        Mesh::HalfedgeHandle h[3];
        Mesh::Point v[3];
        Mesh::Point dir[3];
        for (; h_it != h_end; h_it++, i++)
        {
            h[i] = *h_it;
            auto h1 = m_mesh->from_vertex_handle(h[i]);
            auto h2 = m_mesh->to_vertex_handle(h[i]);

            v[i] = m_mesh->point(h1);
            dir[i] = (m_mesh->point(h2) - v[i]).normalized();
        }

        // 计算余切
        double cot[3];
        for (int i = 0; i < 3; i++)
        {
            int j = (i + 1) % 3;
            double cosv = dir[i] | dir[j];
            cot[i] = 1 / std::tan(std::acos(-cosv));
        }

        // 计算参数坐标
        const auto v10 = v[1] - v[0];
        const auto v10Length = v10.length();
        const auto v20 = v[2] - v[0];

        const auto x = v10 / v10Length;
        const auto z = (x % v20).normalized();
        const auto y = z % x;

        v[0] = Mesh::Point(0, 0, 0);
        v[1] = Mesh::Point(v10Length, 0, 0);
        v[2] = Mesh::Point(v20 | x, v20 | y, 0);

        // 保存信息到半边
        for (int i = 0; i < 3; i++)
        {
            int j = (i + 1) % 3;

            auto dx = v[i] - v[j];
            auto &data = m_mesh->property(m_halfedgeData, h[i]);
            data.cot = cot[j];
            data.dx[0] = dx[0];
            data.dx[1] = dx[1];
        }
    }

    // 顶点编号
    auto &vertices = chart.vertices();
    int size = vertices.size();
    for (int i = 0; i < size; i++)
    {
        auto &v = vertices[i];
        m_mesh->property(m_ids, v) = i;
    }

    // 计算初始矩阵，并保存其预分解形式
    TripletList tripleList;
    for (int i = 0; i < size; i++)
    {
        auto &vi = vertices[i];
        int row = m_mesh->property(m_ids, vi);

        // 查找指向外侧的半边
        auto oh_it = m_mesh->voh_begin(vi), oh_end = m_mesh->voh_end(vi);
        for (; oh_it != oh_end; oh_it++)
        {
            // 跳过边界半边
            if (oh_it->is_boundary())
                continue;

            // 排除不属于图表的半边
            auto face = m_mesh->face_handle(*oh_it);
            if (m_mesh->property(m_chartIds, face) != chart.id())
                continue;

            auto vj = m_mesh->to_vertex_handle(*oh_it);
            int col = m_mesh->property(m_ids, vj);
            auto &data = m_mesh->property(m_halfedgeData, *oh_it);

            tripleList.push_back({row, row, data.cot});
            tripleList.push_back({row, col, -data.cot});
        }

        // 查找指向内侧的半边
        auto ih_it = m_mesh->vih_begin(vi), ih_end = m_mesh->vih_end(vi);
        for (; ih_it != ih_end; ih_it++)
        {
            // 跳过边界半边
            if (ih_it->is_boundary())
                continue;

            // 排除不属于图表的半边
            auto face = m_mesh->face_handle(*ih_it);
            if (m_mesh->property(m_chartIds, face) != chart.id())
                continue;

            auto vj = m_mesh->from_vertex_handle(*ih_it);
            int col = m_mesh->property(m_ids, vj);
            auto &data = m_mesh->property(m_halfedgeData, *ih_it);

            tripleList.push_back({row, row, data.cot});
            tripleList.push_back({row, col, -data.cot});
        }
    }

    std::cout << "Chart vertices num: " << size << std::endl;

    SparseMatrix M(size, size);
    M.setFromTriplets(tripleList.begin(), tripleList.end());

    // 构造加权系数矩阵
    TripletList diag;
    for (int i = 0; i < size; i++)
        diag.push_back({i, i, m_lambda * size});
    SparseMatrix D(size, size);
    D.setFromTriplets(diag.begin(), diag.end());

    SparseMatrix MT = M.transpose();
    SparseMatrix MTM = MT * M + D;
    MTM.makeCompressed();

    m_solver.compute(MTM);
    auto info = m_solver.info();
    switch (info)
    {
    case Eigen::Success:
        std::cout << "Factorize successfully." << std::endl;
        break;
    case Eigen::NumericalIssue:
        std::cout << "Factorize NumericalIssue." << std::endl;
        break;
    case Eigen::NoConvergence:
        std::cout << "Factorize NoConvergence." << std::endl;
        break;
    case Eigen::InvalidInput:
        std::cout << "Factorize InvalidInput." << std::endl;
        break;
    default:
        break;
    }

    return MT;
}

void ARAPBuilder::localPhase(const Chart &chart)
{
    auto &faces = chart.faces();
    for (auto &face : faces)
    {
        MatrixX S(2, 2);
        S.setZero();

        auto h_it = m_mesh->fh_ccwbegin(face), h_end = m_mesh->fh_ccwend(face);
        for (; h_it != h_end; h_it++)
        {
            auto &data = m_mesh->property(m_halfedgeData, *h_it);
            auto v1 = m_mesh->from_vertex_handle(*h_it);
            auto v2 = m_mesh->to_vertex_handle(*h_it);

            MatrixXx1 u(2);
            auto uv = m_mesh->texcoord2D(v1) - m_mesh->texcoord2D(v2);
            u[0] = uv[0];
            u[1] = uv[1];

            S = S + u * data.dx.transpose() * data.cot;
        }

        // 奇异值分解，然后更新 Lt
        Eigen::JacobiSVD<MatrixX> svd(S, Eigen::ComputeFullU | Eigen::ComputeFullV);
        m_mesh->property(m_Lts, face) = svd.matrixU() * svd.matrixV().transpose();
    }
}

void ARAPBuilder::globalPhase(const Chart &chart, const SparseMatrix &M)
{
    auto &vertices = chart.vertices();
    int size = vertices.size();
    MatrixX b(size, 2);
    b.setZero();

    for (int i = 0; i < size; i++)
    {
        auto &vi = vertices[i];
        int row = m_mesh->property(m_ids, vi);

        // 查找指向外侧的半边
        auto oh_it = m_mesh->voh_begin(vi), oh_end = m_mesh->voh_end(vi);
        for (; oh_it != oh_end; oh_it++)
        {
            // 排除掉边界半边
            if (oh_it->is_boundary())
                continue;

            // 排除不属于图表的半边
            auto face = m_mesh->face_handle(*oh_it);
            if (m_mesh->property(m_chartIds, face) != chart.id())
                continue;

            auto vj = m_mesh->to_vertex_handle(*oh_it);
            int col = m_mesh->property(m_ids, vj);
            auto &data = m_mesh->property(m_halfedgeData, *oh_it);

            MatrixX &Lt = m_mesh->property(m_Lts, face);
            MatrixXx1 db = Lt * data.dx * data.cot;

            // 设置系数
            b(row, 0) += db[0];
            b(row, 1) += db[1];
        }

        // 查找指向内侧的半边
        auto ih_it = m_mesh->vih_begin(vi), ih_end = m_mesh->vih_end(vi);
        for (; ih_it != ih_end; ih_it++)
        {
            // 排除掉边界半边
            if (ih_it->is_boundary())
                continue;

            // 排除不属于图表的半边
            auto face = m_mesh->face_handle(*ih_it);
            if (m_mesh->property(m_chartIds, face) != chart.id())
                continue;

            auto vj = m_mesh->from_vertex_handle(*ih_it);
            int col = m_mesh->property(m_ids, vj);
            auto &data = m_mesh->property(m_halfedgeData, *ih_it);

            MatrixX &Lt = m_mesh->property(m_Lts, face);
            MatrixXx1 db = -Lt * data.dx * data.cot;

            // 设置系数
            b(row, 0) += db[0];
            b(row, 1) += db[1];
        }
    }

    // 求解全局系统
    MatrixX x = m_solver.solve(M * b);
    for (int i = 0; i < size; i++)
    {
        auto &vi = vertices[i];
        Mesh::TexCoord2D tex;
        tex[0] = x(i, 0);
        tex[1] = x(i, 1);

        // 更新参数
        m_mesh->set_texcoord2D(vi, tex);
    }
}

void ARAPBuilder::storeUVs(const Chart &chart)
{
    // // 初始化最大最小纹理坐标
    // auto minUV = Mesh::TexCoord2D(MIN_DEFAULT, MIN_DEFAULT);
    // auto maxUV = Mesh::TexCoord2D(MAX_DEFAULT, MAX_DEFAULT);

    // for (const auto &vertex : chart.vertices())
    // {
    //     // 固定顶点已经有参数坐标，只需要考虑活动顶点
    //     auto uv = m_mesh->texcoord2D(vertex);

    //     // 记录最大、最小纹理坐标
    //     minUV.minimize(uv);
    //     maxUV.maximize(uv);
    // }

    // auto diff = maxUV - minUV;
    // const auto length = diff.max();

    // // 规范化纹理坐标
    // for (const auto &vertex : chart.vertices())
    // {
    //     auto uv = m_mesh->texcoord2D(vertex);
    //     const auto nuv = (uv - minUV) / length;

    //     // 记录顶点的纹理坐标
    //     m_mesh->set_texcoord2D(vertex, nuv);
    // }

    // 坐标分离为 2 个随机变量
    int idx = 0;
    int n = chart.vertices().size();
    Eigen::VectorXd X(n), Y(n);
    for (const auto &vertex : chart.vertices())
    {
        // 固定顶点已经有参数坐标，只需要考虑活动顶点
        auto uv = m_mesh->texcoord2D(vertex);
        X[idx] = uv[0];
        Y[idx] = uv[1];
        idx++;
    }

    // 计算期望点
    double E[] = {0, 0};
    E[0] = X.sum() / n;
    E[1] = Y.sum() / n;

    // 计算协方差矩阵
    Eigen::MatrixXd B(2, 2);
    B.setZero();
    for (int i = 0; i < n; i++)
    {
        B(0, 0) += X[i] * X[i];
        B(0, 1) += X[i] * Y[i];

        B(1, 0) += Y[i] * X[i];
        B(1, 1) += Y[i] * Y[i];
    }
    B = B / n;

    // 计算 Xi_i 和 Xi_j 的协方差
    for (int i = 0; i < 2; i++)
        for (int j = 0; j < 2; j++)
            B(i, j) -= E[i] * E[j];

    // 特征分解
    Eigen::EigenSolver<Eigen::MatrixXd> es(B);

    // 获得特征值，取实部
    Eigen::VectorXcd cval = es.eigenvalues();
    Eigen::VectorXd val(cval.size());
    for (int i = 0; i < val.size(); i++)
        val[i] = cval[i].real();

    // 获得特征向量的矩阵，取实部
    Eigen::MatrixXcd cV = es.eigenvectors();
    Eigen::MatrixXd V(cV.rows(), cV.cols());
    for (int i = 0; i < V.rows(); i++)
        for (int j = 0; j < V.cols(); j++)
            V(i, j) = cV(i, j).real();

    int order[] = {0, 1};
    if (val[1] > val[0])
    {
        std::swap(order[0], order[1]);
        std::swap(val[0], val[1]);
    }

    OpenMesh::Vec2d x(V(0, order[0]), V(1, order[0]));
    OpenMesh::Vec2d y(V(0, order[1]), V(1, order[1]));

    idx = 0;
    for (const auto &vertex : chart.vertices())
    {
        auto uv = m_mesh->texcoord2D(vertex);
        OpenMesh::Vec2d coord(uv[0] - E[0], uv[1] - E[1]);
        uv[0] = x.dot(coord);
        uv[1] = y.dot(coord);

        // 记录顶点的纹理坐标
        m_mesh->set_texcoord2D(vertex, uv);
    }
}
